import numpy as np
import matplotlib.pyplot as plt  # 绘图核心库
import torch
from torch.utils.data import DataLoader, Dataset

# -------------------------- 1. 自定义数据集（不变） --------------------------
class DiabetesDataset(Dataset):
    def __init__(self, file_path):
        xy = np.loadtxt(file_path, delimiter=",", dtype=np.float32)
        self.len = xy.shape[0]
        self.x_data = torch.from_numpy(xy[:, :-1])  # 8个输入特征
        self.y_data = torch.from_numpy(xy[:, [-1]]) # 1个标签（0/1）

    def __getitem__(self, index):
        return self.x_data[index], self.y_data[index]

    def __len__(self):
        return self.len

# -------------------------- 2. 定义4层全连接模型（不变） --------------------------
class Model(torch.nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.linear1 = torch.nn.Linear(8, 6)
        self.linear2 = torch.nn.Linear(6, 4)
        self.linear3 = torch.nn.Linear(4, 2)
        self.linear4 = torch.nn.Linear(2, 1)
        self.activation = torch.nn.Sigmoid()  # 复用激活函数

    def forward(self, x):
        x = self.activation(self.linear1(x))
        x = self.activation(self.linear2(x))
        x = self.activation(self.linear3(x))
        x = self.activation(self.linear4(x))
        return x

# -------------------------- 3. 主逻辑（新增损失记录+绘图） --------------------------
if __name__ == '__main__':
    # 初始化数据集和DataLoader（移到主逻辑内，避免多进程冲突）
    dataset = DiabetesDataset("diabetes.csv")  # 确保文件路径正确！
    train_loader = DataLoader(
        dataset=dataset,
        batch_size=32,
        shuffle=True,
        num_workers=2
    )

    # 初始化模型、损失函数、优化器
    model = Model()
    criterion = torch.nn.BCELoss(reduction='mean')  # 替换过时参数，消除警告
    optimizer = torch.optim.SGD(model.parameters(), lr=0.1)

    # -------------------------- 关键：初始化列表记录损失 --------------------------
    epoch_list = []  # 存储每一轮的序号（x轴数据）
    loss_list = []   # 存储每一轮的平均损失（y轴数据）

    # 训练循环（100轮）
    for epoch in range(100):
        # 每轮初始化“批次损失总和”和“批次数量”，用于计算该轮平均损失
        batch_loss_sum = 0.0  # 累加当前轮所有批次的损失
        batch_count = 0       # 记录当前轮的批次数

        # 按批次遍历数据
        for i, data in enumerate(train_loader, 0):
            inputs, labels = data

            # 前向传播+计算损失
            y_pred = model(inputs)
            loss = criterion(y_pred, labels)

            # 反向传播+参数更新
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            # 累加当前批次损失，统计批次数
            batch_loss_sum += loss.item()  # .item()转为Python数值，避免张量累积
            batch_count += 1

        # -------------------------- 关键：计算当前轮的平均损失 --------------------------
        epoch_avg_loss = batch_loss_sum / batch_count  # 平均损失 = 总损失 / 批次数
        epoch_list.append(epoch)                       # 记录轮次
        loss_list.append(epoch_avg_loss)               # 记录该轮平均损失

        # 打印每轮的平均损失（更易观察收敛趋势）
        print(f"epoch: {epoch:3d}, 平均损失: {epoch_avg_loss:.6f}")

    # -------------------------- 关键：绘制loss-epoch图 --------------------------
    plt.figure(figsize=(10, 6))  # 设置图的尺寸（宽10，高6）
    plt.plot(epoch_list, loss_list, color='#1f77b4', linewidth=2, marker='o', markersize=4)
    # 说明：
    # color='#1f77b4'：蓝色线条（matplotlib默认经典色）
    # linewidth=2：线条粗细
    # marker='o'：每个数据点用圆形标记
    # markersize=4：标记大小

    # 设置图表标签和标题
    plt.xlabel('Epoch (训练轮次)', fontsize=12)  # x轴标签
    plt.ylabel('Average Loss (平均损失)', fontsize=12)  # y轴标签
    plt.title('Diabetes Classification: Loss vs Epoch', fontsize=14, pad=20)  # 标题

    # 设置网格（便于读取数值）
    plt.grid(True, linestyle='--', alpha=0.7)  # linestyle='--'：虚线网格；alpha=0.7：透明度

    # 保存图片（可选，建议保存为高清格式）
    plt.savefig('loss_vs_epoch.png', dpi=300, bbox_inches='tight')  # dpi=300：高清；bbox_inches='tight'：避免标签被截断

    # 显示图片
    plt.show()